#include "controller_tracker/controller_tracker_node.hpp"
#include <ament_index_cpp/get_package_share_directory.hpp>

namespace controller_tracker {
/*
    * 1. 创建一个订阅器，订阅AprilTagDetectionArray类型的消息
    * 2. 实现tagCallback函数，处理订阅到的消息
    * 3. 在tagCallback函数中，使用tf2_ros::Buffer和tf2_ros::TransformListener来查询apriltag的位姿
    * 4. 在tagCallback函数中，处理apriltag的位姿，实现控制器的跟踪功能
    */
ControllerTrackerNode::ControllerTrackerNode(const rclcpp::NodeOptions & options)
: Node("controller_tracker_node",options)
{
    RCLCPP_INFO(this->get_logger(), "Starting TrackerNode!");

    auto callback = [this](const AprilTagDetectionArray::SharedPtr msg) {
        this->tagCallback(msg);
    };

    std::string package_share_dir = ament_index_cpp::get_package_share_directory("controller_tracker");
    // 加载urdf文件
    urdf_path_ = this->declare_parameter("target_urdf_path", package_share_dir + "/urdf/controller_target.urdf.xacro");
    // 
    if (!model_.initFile(urdf_path_)) { // 初始化urdf模型
        RCLCPP_ERROR(this->get_logger(), "Failed to parse urdf file");
        return;
    } 

    // 加载apriltag的位姿
    tag_ids_ = this->declare_parameter<std::vector<int64_t>>("tags_positions.ids", std::vector<int64_t>{});
    auto positions_x = this->declare_parameter<std::vector<double>>("tags_positions.positions_x", std::vector<double>{});
    auto positions_y = this->declare_parameter<std::vector<double>>("tags_positions.positions_y", std::vector<double>{});
    auto positions_z = this->declare_parameter<std::vector<double>>("tags_positions.positions_z", std::vector<double>{});
    auto orientations_x = this->declare_parameter<std::vector<double>>("tags_positions.orientations_x", std::vector<double>{});
    auto orientations_y = this->declare_parameter<std::vector<double>>("tags_positions.orientations_y", std::vector<double>{});
    auto orientations_z = this->declare_parameter<std::vector<double>>("tags_positions.orientations_z", std::vector<double>{});
    auto orientations_w = this->declare_parameter<std::vector<double>>("tags_positions.orientations_w", std::vector<double>{});

    // 检查参数是否加载成功
    if (tag_ids_.empty() || positions_x.empty() || positions_y.empty() || positions_z.empty() ||
        orientations_x.empty() || orientations_y.empty() || orientations_z.empty() || orientations_w.empty()) {
        RCLCPP_ERROR(this->get_logger(), "Failed to load tag positions from parameters");
        throw std::runtime_error("Failed to load tag positions from parameters");
    }

    // 检查所有数组大小是否一致
    size_t size = tag_ids_.size();
    if (positions_x.size() != size || positions_y.size() != size || positions_z.size() != size ||
        orientations_x.size() != size || orientations_y.size() != size || 
        orientations_z.size() != size || orientations_w.size() != size) {
        RCLCPP_ERROR(this->get_logger(), "Mismatch in tag parameters array sizes");
        throw std::runtime_error("Mismatch in tag parameters array sizes");
    }

    // 构建positions和orientations向量
    positions_.resize(size);
    orientations_.resize(size);
    for (size_t i = 0; i < size; ++i) {
        positions_[i] = {positions_x[i], positions_y[i], positions_z[i]};
        orientations_[i] = {orientations_x[i], orientations_y[i], orientations_z[i], orientations_w[i]};
    }

    RCLCPP_INFO(this->get_logger(), "Successfully loaded %ld tag positions", size);


    tf_buffer_ = std::make_shared<tf2_ros::Buffer>(this->get_clock());
    tf_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);


    sub_ = this->create_subscription<AprilTagDetectionArray>("detections", 10, callback);   // 创建一个订阅器，订阅AprilTagDetectionArray类型的消息
    
    // 创建target_pose发布器
    target_pose_pub_ = this->create_publisher<geometry_msgs::msg::PoseStamped>("target_pose", 10);
}


void ControllerTrackerNode::tagCallback(const AprilTagDetectionArray::SharedPtr msg) {
    double max_decision_margin = -1.0;
    apriltag_msgs::msg::AprilTagDetection best_detection;

    for (const auto& detection : msg->detections) { // 遍历所有检测到的apriltag
        if (detection.id >= 0 && detection.id <= 4 && detection.decision_margin > max_decision_margin) {
            max_decision_margin = detection.decision_margin;
            best_detection = detection;
        }
    }

    if (max_decision_margin > -1.0) { // 如果找到了最佳的apriltag
        std::string tag_frame = "tag36h11:" + std::to_string(best_detection.id);
        RCLCPP_INFO(this->get_logger(), "Best detection: %s", tag_frame.c_str());
        try {
            // 获取变换
            geometry_msgs::msg::TransformStamped transform_stamped = 
            (*tf_buffer_).lookupTransform("camera_optical_frame", tag_frame, rclcpp::Time(0));

            auto tag_position = positions_[best_detection.id];
            auto tag_orientation = orientations_[best_detection.id];

            // 计算target_link的方向
            tf2::Quaternion tag_orientation_quat(tag_orientation[0], tag_orientation[1], tag_orientation[2], tag_orientation[3]);
            tf2::Quaternion transform_stamped_quat(transform_stamped.transform.rotation.x, transform_stamped.transform.rotation.y, transform_stamped.transform.rotation.z, transform_stamped.transform.rotation.w);

            tf2::Quaternion target_orientation_quat = transform_stamped_quat * tag_orientation_quat;

            // 计算target_link的位姿
            geometry_msgs::msg::TransformStamped target_transform;
            target_transform.header.frame_id = "camera_optical_frame";
            target_transform.child_frame_id = "target_link";
            target_transform.transform.translation.x = transform_stamped.transform.translation.x + tag_position[0];
            target_transform.transform.translation.y = transform_stamped.transform.translation.y + tag_position[1];
            target_transform.transform.translation.z = transform_stamped.transform.translation.z + tag_position[2];
            target_transform.transform.rotation.x = target_orientation_quat.x();
            target_transform.transform.rotation.y = target_orientation_quat.y();
            target_transform.transform.rotation.z = target_orientation_quat.z();
            target_transform.transform.rotation.w = target_orientation_quat.w();

            // 创建一个TransformBroadcaster
            static tf2_ros::TransformBroadcaster tf_broadcaster(this);

            // 发布变换
            tf_broadcaster.sendTransform(target_transform);

            // 发布target_pose
            geometry_msgs::msg::PoseStamped target_pose;
            target_pose.header = target_transform.header;
            target_pose.pose.position.x = target_transform.transform.translation.x;
            target_pose.pose.position.y = target_transform.transform.translation.y;
            target_pose.pose.position.z = target_transform.transform.translation.z;
            target_pose.pose.orientation = target_transform.transform.rotation;

            target_pose_pub_->publish(target_pose);

        } catch (tf2::TransformException &ex) {
            RCLCPP_WARN(this->get_logger(), "%s", ex.what());
        }
    }
}

}  // namespace controller_tracker

#include "rclcpp_components/register_node_macro.hpp"

// Register the component with class_loader.
// This acts as a sort of entry point, allowing the component to be discoverable when its library
// is being loaded into a running process.
RCLCPP_COMPONENTS_REGISTER_NODE(controller_tracker::ControllerTrackerNode)
